{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}

-- | Module: Data.Vector.Mutable.Sized.Sort
-- Copyright: (c) 2018 Koz Ross
-- Maintainer: Koz Ross <koz.ross@retro-freedom.nz>
-- Stability: Experimental
-- Portability: Non-portable
--
-- A range of implementations of sorting algorithms, based on the
-- vector-algorithms package. These sort in-place.
module Data.Vector.Mutable.Sized.Sort
  ( Lexicographic
  , Comparison
  , americanFlagSort
  , americanFlagSortBy
  , heapSort
  , heapSortBy
  , heapSortOn
  , insertionSort
  , insertionSortBy
  , insertionSortOn
  , introSort
  , introSortBy
  , introSortOn
  , mergeSort
  , mergeSortBy
  , mergeSortOn
  , timSort
  , timSortBy
  , timSortOn
  , discriminantSort
  , discriminantSortOn
  ) where

import Control.Monad.Primitive
import Data.Coerce
import Data.Discrimination
import Data.Foldable
import Data.Function
import Data.Functor.Contravariant
import Data.Vector.Algorithms.AmericanFlag (Lexicographic)

import qualified Data.Vector.Algorithms.AmericanFlag as AF
import qualified Data.Vector.Algorithms.Heap as H
import qualified Data.Vector.Algorithms.Insertion as I
import qualified Data.Vector.Algorithms.Intro as IS
import qualified Data.Vector.Algorithms.Merge as M
import qualified Data.Vector.Algorithms.Tim as T
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Generic.Mutable.Sized as VGMS

-- | Sorts using American flag sort. For details, see
-- Data.Vector.Algorithms.AmericanFlag.
--
-- __Complexity:__ \(O(n)\)
{-# INLINE americanFlagSort #-}
americanFlagSort ::
     (Lexicographic a, Ord a, PrimMonad m, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> m ()
americanFlagSort = AF.sort . VGMS.fromSized

-- | Fine-grained version of 'americanFlagSort'.
--
-- __Complexity:__ \(O(n)\)
{-# INLINE americanFlagSortBy #-}
americanFlagSortBy ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a -- ^ Comparison for insertion sort fallback.
  -> (a -> Int -> Bool) -- ^ Is a stripe complete?
  -> Int -- ^ Number of buckets necessary.
  -> (Int -> a -> Int) -- ^ Big-endian radix function.
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
americanFlagSortBy comp strip buckets radix =
  AF.sortBy (coerce comp) strip buckets radix . VGMS.fromSized

-- | Sorts using heapsort with a quarternary heap. For details, see
-- Data.Vector.Algorithms.Heap.
--
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE heapSort #-}
heapSort ::
     (Ord a, PrimMonad m, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> m ()
heapSort = H.sort . VGMS.fromSized

-- | As 'heapSort', but with a custom 'Comparison'.
{-# INLINE heapSortBy #-}
heapSortBy ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
heapSortBy comp = H.sortBy (coerce comp) . VGMS.fromSized

-- | As 'heapSort', using a \'projection\' to determine ordering.
{-# INLINE heapSortOn #-}
heapSortOn ::
     (PrimMonad m, Ord b, VGM.MVector v a)
  => (a -> b)
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
heapSortOn f = H.sortBy go . VGMS.fromSized
  where
    go = coerce (compare `on` f)

-- | Sorts using insertion sort. For details, see
-- Data.Vector.Algorithms.Insertion.
--
-- __Complexity:__ \(O(n^{2})\)
{-# INLINE insertionSort #-}
insertionSort ::
     (Ord a, PrimMonad m, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> m ()
insertionSort = I.sort . VGMS.fromSized

-- | As 'insertionSort', but with a custom 'Comparison'.
{-# INLINE insertionSortBy #-}
insertionSortBy ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
insertionSortBy comp = I.sortBy (coerce comp) . VGMS.fromSized

-- | As 'insertionSort', using a \'projection\' to determine ordering.
{-# INLINE insertionSortOn #-}
insertionSortOn ::
     (PrimMonad m, Ord b, VGM.MVector v a)
  => (a -> b)
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
insertionSortOn f = I.sortBy go . VGMS.fromSized
  where
    go = coerce (compare `on` f)

-- | Sorts using introsort. For details, see Data.Vector.Algorithms.Intro.
--
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE introSort #-}
introSort ::
     (Ord a, PrimMonad m, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> m ()
introSort = IS.sort . VGMS.fromSized

-- | As 'introSort', but with a custom 'Comparison'.
{-# INLINE introSortBy #-}
introSortBy ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
introSortBy comp = IS.sortBy (coerce comp) . VGMS.fromSized

-- | As 'introSort', using a \'projection\' to determine ordering.
{-# INLINE introSortOn #-}
introSortOn ::
     (PrimMonad m, Ord b, VGM.MVector v a)
  => (a -> b)
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
introSortOn f = IS.sortBy go . VGMS.fromSized
  where
    go = coerce (compare `on` f)

-- | Sorts using top-down mergesort. For details, see
-- Data.Vector.Algorithms.Merge.
--
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE mergeSort #-}
mergeSort ::
     (Ord a, PrimMonad m, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> m ()
mergeSort = M.sort . VGMS.fromSized

-- | As 'mergeSort', but with a custom 'Comparison'.
{-# INLINE mergeSortBy #-}
mergeSortBy ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
mergeSortBy comp = M.sortBy (coerce comp) . VGMS.fromSized

-- | As 'mergeSort', using a \'projection\' to determine ordering.
{-# INLINE mergeSortOn #-}
mergeSortOn ::
     (PrimMonad m, Ord b, VGM.MVector v a)
  => (a -> b)
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
mergeSortOn f = M.sortBy go . VGMS.fromSized
  where
    go = coerce (compare `on` f)

-- | Sorts using Timsort. For details, see Data.Vector.Algorithms.Tim.
--
-- __Complexity:__ \(O(n \log(n))\)
{-# INLINE timSort #-}
timSort ::
     (Ord a, PrimMonad m, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> m ()
timSort = T.sort . VGMS.fromSized

-- | As 'timSort', but with a custom 'Comparison'.
{-# INLINE timSortBy #-}
timSortBy ::
     (PrimMonad m, VGM.MVector v a)
  => Comparison a
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
timSortBy comp = T.sortBy (coerce comp) . VGMS.fromSized

-- | As 'timSort', using a \'projection\' to determine ordering.
{-# INLINE timSortOn #-}
timSortOn ::
     (PrimMonad m, Ord b, VGM.MVector v a)
  => (a -> b)
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
timSortOn f = T.sortBy go . VGMS.fromSized
  where
    go = coerce (compare `on` f)

-- | Sorts using a stable, ordered discriminator. For details, see
-- 'Data.Discrimination'.
--
-- __Complexity:__ \(O(n)\)
{-# INLINE discriminantSort #-}
discriminantSort ::
     (PrimMonad m, Sorting a, VGM.MVector v a)
  => VGMS.MVector v n (PrimState m) a
  -> m ()
discriminantSort = overMutable sort

-- | As 'discriminantSort', using a \'projection\' to determine the discriminant
-- used.
{-# INLINE discriminantSortOn #-}
discriminantSortOn ::
     (PrimMonad m, Sorting b, VGM.MVector v a)
  => (a -> b)
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
discriminantSortOn f = overMutable (sortWith f)

-- helper for 'injecting' into a mutable unsized underlying vector
{-# INLINE overMutable #-}
overMutable ::
     (PrimMonad m, VGM.MVector v a)
  => ([a] -> [a])
  -> VGMS.MVector v n (PrimState m) a
  -> m ()
overMutable f vec = do
  let v = VGMS.fromSized vec
  let ixes = [0 .. VGM.length v - 1]
  ell <- traverse (VGM.unsafeRead v) ixes
  let ell' = f ell
  traverse_ (uncurry (VGM.unsafeWrite v)) (zip ixes ell')
